//
// Created by jsjxsy on 2024/8/1.
//
#include <stdio.h>
#include <libavutil/log.h>
#include <stdlib.h>
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libswscale/swscale.h>
#define WORD uint16_t
#define DWORD uint32_t
#define LONG int32_t

typedef struct tagBITMAPFILEHEADER {
  WORD bfType;
  DWORD bfSize;
  WORD bfReserved1;
  WORD bfReserved2;
  DWORD bfOffBits;
} BITMAPFILEHEADER, *LPBITMAPFILEHEADER, *PBITMAPFILEHEADER;

typedef struct tagBITMAPINFOHEADER {
  DWORD biSize;
  LONG biWidth;
  LONG biHeight;
  WORD biPlanes;
  WORD biBitCount;
  DWORD biCompression;
  DWORD biSizeImage;
  LONG biXPelsPerMeter;
  LONG biYPelsPerMeter;
  DWORD biClrUsed;
  DWORD biClrImportant;
} BITMAPINFOHEADER, *LPBITMAPINFOHEADER, *PBITMAPINFOHEADER;

static void saveBMP(struct SwsContext *swsCtx, AVFrame *frame, int width, int height, char *name) {
  int dataSize = width * height * 3;
  FILE *f = NULL;
  AVFrame *frameBGR = av_frame_alloc();
  frameBGR->width = width;
  frameBGR->height = height;
  frameBGR->format = AV_PIX_FMT_BGR24;
  av_frame_get_buffer(frameBGR, 0);
  sws_scale(swsCtx,
            (const uint8_t *const *) frame->data,
            frame->linesize,
            0,
            frame->height,
            frameBGR->data,
            frameBGR->linesize);
  BITMAPINFOHEADER infoHeader;
  infoHeader.biSize = sizeof(BITMAPINFOHEADER);
  infoHeader.biWidth = width;
  infoHeader.biHeight = height * (-1);
  infoHeader.biBitCount = 24;
  infoHeader.biCompression = 0;
  infoHeader.biSizeImage = 0;
  infoHeader.biClrImportant = 0;
  infoHeader.biClrUsed = 0;
  infoHeader.biXPelsPerMeter = 0;
  infoHeader.biYPelsPerMeter = 0;
  infoHeader.biPlanes = 0;
  BITMAPFILEHEADER fileHeader;
  fileHeader.bfType = 0x4d42;
  fileHeader.bfSize = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER) + dataSize;
  fileHeader.bfOffBits = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER);
  f = fopen(name, "wb");
  fwrite(&fileHeader, sizeof(BITMAPFILEHEADER), 1, f);
  fwrite(&infoHeader, sizeof(BITMAPINFOHEADER), 1, f);
  fwrite(frameBGR->data[0], 1, dataSize, f);
  fclose(f);
  av_freep(&frameBGR->data[0]);
  av_freep(frameBGR);
}

static void savePic(unsigned char *buf, int lineSize, int width, int height, char *name) {
  FILE *f;
  f = fopen(name, "wb");
  fprintf(f, "P5\n%d %d\n%d\n", width, height, 255);
  for (int i = 0; i < height; i++) {
    fwrite(buf + i * lineSize, 1, width, f);
  }
  fclose(f);
}

static int decode(AVCodecContext *ctx, struct SwsContext *swsCtx, AVFrame *frame, AVPacket *packet, char *dist) {
  int ret = -1;
  char buf[1024];
  ret = avcodec_send_packet(ctx, packet);
  if (ret < 0) {
    av_log(NULL, AV_LOG_ERROR, "failed to send packet to decoder: %s \n", av_err2str(ret));
    goto _END;
  }

  while (ret >= 0) {
    ret = avcodec_receive_frame(ctx, frame);
    if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
      return 0;
    } else if (ret < 0) {
      return -1;
    }
    snprintf(buf, sizeof(buf), "%s-%lld.bmp", dist, ctx->frame_num);
//    savePic(frame->data[0], frame->linesize[0], frame->width, frame->height, buf);
    saveBMP(swsCtx, frame, frame->width, frame->height, buf);
    if (packet) {
      av_packet_unref(packet);
    }
  }
  _END:
  return 0;
}

int main(int argc, char *argv[]) {
  char *src;
  char *dist;
  AVFormatContext *pFmtCtx = NULL;
  AVStream *inStream = NULL;
  const AVCodec *codec = NULL;
  AVCodecContext *ctx = NULL;
  AVPacket *pkt = NULL;
  AVFrame *frame = NULL;
  int ret = -1;
  int index = -1;
  struct SwsContext *swsCtx = NULL;
  //1
  av_log_set_level(AV_LOG_DEBUG);
  if (argc < 3) {
    av_log(NULL, AV_LOG_INFO, "arguments must be more than 3\n");
    exit(-1);
  }
  src = argv[1];
  dist = argv[2];
  //2
  ret = avformat_open_input(&pFmtCtx, src, NULL, NULL);
  if (ret < 0) {
    av_log(NULL, AV_LOG_ERROR, "%s\n", av_err2str(ret));
    exit(-1);
  }
  //3
  index = av_find_best_stream(pFmtCtx, AVMEDIA_TYPE_VIDEO, -1, -1, NULL, 0);
  if (index < 0) {
    av_log(pFmtCtx, AV_LOG_ERROR, "does not include video stream \n");
    goto _ERROR;
  }
  inStream = pFmtCtx->streams[index];
  //4
  codec = avcodec_find_decoder(inStream->codecpar->codec_id);
  if (!codec) {
    av_log(NULL, AV_LOG_ERROR, "don't find Codec: libx264\n");
    exit(-1);
  }
  //5
  ctx = avcodec_alloc_context3(codec);
  if (!ctx) {
    av_log(NULL, AV_LOG_ERROR, "no memory!\n");
    goto _ERROR;
  }
  ret = avcodec_parameters_to_context(ctx, inStream->codecpar);
  if (ret < 0) {
    av_log(NULL, AV_LOG_ERROR, "could not copy codecpar to codec ctx!\n");
    goto _ERROR;
  }
  //6
  ret = avcodec_open2(ctx, codec, NULL);
  if (ret < 0) {
    av_log(NULL, AV_LOG_ERROR, "don't open codec %s\n", av_err2str(ret));
    goto _ERROR;
  }
  //6.1
  swsCtx = sws_getCachedContext(NULL,
                                ctx->width,
                                ctx->height,
                                AV_PIX_FMT_YUV420P,//AV_PIX_FMT_YUV420P,
                                640,
                                360,
                                AV_PIX_FMT_BGR24,
                                SWS_BICUBIC,
                                NULL,
                                NULL,
                                NULL);

  //7
  frame = av_frame_alloc();
  if (!frame) {
    av_log(NULL, AV_LOG_ERROR, "no memory\n");
    goto _ERROR;
  }
  //8
  pkt = av_packet_alloc();
  if (!pkt) {
    av_log(NULL, AV_LOG_ERROR, "no memory\n");
    goto _ERROR;
  }

  //9
  while (av_read_frame(pFmtCtx, pkt) >= 0) {
    if (pkt->stream_index == index) {
      decode(ctx, swsCtx, frame, pkt, dist);
    }
  }
  decode(ctx, swsCtx, frame, NULL, dist);
  //10
  _ERROR:
  if (pFmtCtx) {
    avformat_close_input(&pFmtCtx);
    pFmtCtx = NULL;
  }

  if (ctx) {
    avcodec_free_context(&ctx);
    ctx = NULL;
  }

  if (frame) {
    av_frame_free(&frame);
    frame = NULL;
  }

  if (pkt) {
    av_packet_free(&pkt);
    pkt = NULL;
  }

  printf("gen pic complete");
  return 0;
}

